#!/usr/bin/python3
# -*- coding: utf-8 -*-

from collections import Iterable
from os.path import isfile

from paramiko import SSHClient, Transport, WarningPolicy
from paramiko.config import SSH_PORT

from ikanda.core.callbacks import CallbackLogger
from ikanda.core.command import Environment, Resource, UploadCommand, ExecuteCommand
from ikanda.core.credentials import Credential, PasswordCredential, KeyPairCredential
from ikanda.core.envnames import *
from ikanda.core.validators import reject_if_null, reject_if_blank, reject_if_non_positive

_DEFAULT_AUTH_TIMEOUT = 10


def _client_name(client: SSHClient) -> str:
    """获取表示SSH客户端的名字。

    :param client: SSH客户端。
    """
    reject_if_null(client)

    t = client.get_transport()
    if isinstance(t, Transport):
        host, port = t.getpeername()
        username = t.get_username()
        return '{username}@{host}:{port}'.format(username=username, host=host, port=port)
    else:
        return '<null>'


class SshResource(Resource):
    """用于调用SSH的底层资源。"""

    __slots__ = ('__clients',)

    def __init__(self):
        super().__init__()

        self.__clients = dict()

    def close(
            self,
            callback: CallbackLogger
    ):
        # 依次关闭所有已创建的SSH客户端。
        for key, client in self.__clients.items():
            if client:
                callback.info('[{client_name}] exit', client_name=_client_name(client))
                client.close()

    def get(
            self,
            hostname: str,
            port: int,
            credential: Credential,
            auth_timeout: int,
            callback: CallbackLogger
    ) -> SSHClient:
        """ 根据key获取或者创建SSH客户端。

        :param hostname: 主机名。
        :param port: 端口号。
        :param credential: 登录凭据。
        :param auth_timeout: 登录超时时间。
        :param callback: 用于记录日志的回调类。
        """
        port = int(port) if port else SSH_PORT
        auth_timeout = int(auth_timeout) if auth_timeout else _DEFAULT_AUTH_TIMEOUT

        reject_if_blank(hostname)
        reject_if_non_positive(port)
        reject_if_null(credential)
        reject_if_non_positive(auth_timeout)

        key = (credential.username, hostname, port)
        result = self.__clients.get(key, None)
        if not result:
            result = SSHClient()
            result.set_missing_host_key_policy(WarningPolicy)
            callback.info('ssh {username}@{hostname}:{port}', username=credential.username, hostname=hostname,
                          port=port)
            if isinstance(credential, PasswordCredential):
                result.connect(hostname=hostname, port=port, username=credential.username,
                               password=credential.password, auth_timeout=auth_timeout)
            elif isinstance(credential, KeyPairCredential):
                result.connect(hostname=hostname, port=port, username=credential.username,
                               pkey=credential.pkey, auth_timeout=auth_timeout)
            else:
                raise ValueError('unknown credential: ' + str(credential))
            self.__clients[key] = result

        return result


class SshUploadCommand(UploadCommand):
    """使用SSH协议上传文件。"""

    def __init__(self,
                 src: str,
                 dst: str,
                 description: str = '',
                 check_required: bool = False
                 ):
        super().__init__(src, dst, description, check_required, node_dep=True)

    def __call__(
            self,
            environment: Environment,
            resources: list,
            callback: CallbackLogger
    ):
        src = self.abspath(self.interpolate(self.src, environment), environment)
        dst = self.interpolate(self.dst, environment)

        if not isfile(src):
            callback.error('cannot find file {src}', src=src)
            return

        def on_put(
                sent,
                total
        ):
            callback.progress(sent * 100 / total)

        ssh_res = self.get_resource(SshResource, resources)
        if isinstance(ssh_res, SshResource):
            hostname = environment.get(NODE_HOSTNAME)
            if not isinstance(hostname, str) or hostname.isspace():
                raise RuntimeError('Illegal hostname: ' + repr(hostname))
            port = environment.getint(NODE_SSH_PORT, 22)
            if port <= 0:
                raise RuntimeError('Illegal ssh port: ' + repr(port))
            credential = environment.get(NODE_SSH_CREDENTIAL, None)
            auth_timeout = environment.getint(NODE_SSH_VERIFY_TIMEOUT, 0)
            client = ssh_res.get(hostname, port, credential, auth_timeout, callback)
            callback.info('[{client_name}] put {src} to {dst}', client_name=_client_name(client), src=src, dst=dst)
            with client.open_sftp() as sftp:
                sftp.put(src, dst, confirm=self.check_required, callback=on_put)
        else:
            raise ValueError('cannot find ' + str(SshResource))


class SshExecuteCommand(ExecuteCommand):
    """使用SSH协议在远程主机上执行程序。"""

    def __init__(self,
                 command_lines: (str, Iterable),
                 redirect_encoding: str = None,
                 timeout: int = None,
                 description: str = '',
                 check_required: bool = False
                 ):
        super().__init__(command_lines=command_lines, redirect_encoding=redirect_encoding, timeout=timeout,
                         description=description, check_required=check_required, node_dep=True)

    def call_single(
            self,
            line: str,
            environment: Environment,
            resources: list,
            callback: CallbackLogger
    ):
        reject_if_blank(line)

        ssh_res = self.get_resource(SshResource, resources)
        if isinstance(ssh_res, SshResource):
            hostname = environment.getstr(NODE_HOSTNAME).strip()
            if not hostname:
                raise RuntimeError('Illegal hostname: ' + repr(hostname))
            port = environment.getint(NODE_SSH_PORT, 22)
            if port <= 0:
                raise RuntimeError('Illegal ssh port: ' + repr(port))
            credential = environment.get(NODE_SSH_CREDENTIAL, None)
            auth_timeout = environment.getint(NODE_SSH_VERIFY_TIMEOUT, 0)
            client = ssh_res.get(hostname, port, credential, auth_timeout, callback)
            callback.info('[{client_name}] {line}', client_name=_client_name(client), line=line)
            stdin_, stdout_, stderr_ = client.exec_command(line, timeout=self.timeout, get_pty=True)
            for t in stdout_.readlines():
                callback.verbose('[{client_name}] >> {t}', client_name=_client_name(client), t=t.strip())
            for t in stderr_.readlines():
                callback.verbose('[{client_name}] >> {t}', client_name=_client_name(client), t=t.strip())
            return 0
        else:
            raise ValueError('cannot find ' + str(SshResource))
